/**
 * @Author: dn-jinmin
 * @File:  routerchain
 * @Version: 1.0.0
 * @Date: 2024/5/15
 * @Description:
 */

package chainx

import (
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/prompts"
)

// Prompt for the router chain in the multi-prompt chain.
const (
	_destinations = "destinations"
	_input        = "input"
	_text         = "text"
	_formatting   = "formatting"

	MULTI_PROMPT_ROUTER_TEMPLATE = `Given a raw text input to a language model select the model prompt best suited for 
the input. You will be given the names of the available prompts and a description of 
what the prompt is best suited for. You may also revise the original input if you 
think that revising it will ultimately lead to a better response from the language 
model.

<< FORMATTING >>
Return a markdown code snippet with a JSON object formatted to look like:
{{.formatting}}

REMEMBER: "destination" MUST be one of the candidate prompt names specified below OR \
it can be "DEFAULT" if the input is not well suited for any of the candidate prompts.
REMEMBER: "next_inputs" can just be the original input if you don't think any \
modifications are needed.

<< CANDIDATE PROMPTS >>
{{.destinations}}

<< INPUT >>
{{.input}}
`
)

var (
	_prompt = prompts.NewPromptTemplate(MULTI_PROMPT_ROUTER_TEMPLATE, []string{_destinations,
		_input, _formatting})
	_outputParser = outputparser.NewStructured([]outputparser.ResponseSchema{
		{
			Name:        _destinations,
			Description: "name of the prompt to use or \"DEFAULT\"",
		},
		{
			Name:        _input,
			Description: "a potentially modified version of the original input",
		},
	})
)
